import pandas as pd
import numpy as np
from sklearn.tree import DecisionTreeClassifier, plot_tree, export_text
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
import matplotlib.pyplot as plt
import matplotlib
import seaborn as sns

# 设置中文字体支持
matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
matplotlib.rcParams['axes.unicode_minus'] = False

# 数据预处理
def load_data():
    base = pd.read_csv('customer_base.csv')
    behavior = pd.read_csv('customer_behavior_assets.csv')
    
    # 获取最新的客户资产数据（按统计月份排序，取每个客户的最新记录）
    latest_data = behavior.sort_values('stat_month').groupby('customer_id').tail(1)
    
    # 合并基础信息和行为资产数据
    merged = pd.merge(base, latest_data, on='customer_id', how='inner')
    
    # 计算客户未来3个月资产达到100万+的概率
    # 假设资产线性增长，基于最近两个月的数据计算增长率
    customer_sorted = behavior.sort_values(['customer_id', 'stat_month'])
    customer_sorted['prev_assets'] = customer_sorted.groupby('customer_id')['total_assets'].shift(1)
    customer_sorted['asset_growth'] = customer_sorted['total_assets'] - customer_sorted['prev_assets']
    customer_sorted['growth_rate'] = customer_sorted['asset_growth'] / customer_sorted['prev_assets']

    # 计算每个客户的平均增长率
    growth_rates = customer_sorted.groupby('customer_id')['growth_rate'].mean().reset_index()
    growth_rates.columns = ['customer_id', 'avg_growth_rate']

    # 合并增长率数据
    merged = pd.merge(merged, growth_rates, on='customer_id', how='left')

    # 填充缺失的增长率（对于只有一条记录的客户）
    merged['avg_growth_rate'] = merged['avg_growth_rate'].fillna(0)

    # 计算3个月后预测资产
    merged['predicted_assets_3m'] = merged['total_assets'] * (1 + merged['avg_growth_rate']) ** 3

    # 创建目标变量：3个月内资产是否能达到100万+
    merged['target'] = (merged['predicted_assets_3m'] >= 1000000).astype(int)
    
    return merged

# 特征工程 - 处理分类变量
def feature_engineering(data):
    # 数值型特征
    numeric_features = ['age', 'monthly_income', 'total_assets', 'deposit_balance', 'financial_balance', 
                       'fund_balance', 'insurance_balance', 'product_count', 'financial_repurchase_count',
                       'credit_card_monthly_expense', 'investment_monthly_count', 'app_login_count',
                       'app_financial_view_time', 'app_product_compare_count', 'avg_growth_rate']
    
    # 分类型特征
    categorical_features = ['gender', 'occupation', 'occupation_type', 'lifecycle_stage', 
                           'marriage_status', 'city_level', 'branch_name', 'asset_level',
                           'deposit_flag', 'financial_flag', 'fund_flag', 'insurance_flag']
    
    # 处理数值型特征
    X_numeric = data[numeric_features].fillna(0)
    
    # 处理分类型特征 - 使用标签编码
    X_categorical = data[categorical_features].copy()
    label_encoders = {}
    for feature in categorical_features:
        le = LabelEncoder()
        X_categorical[feature] = le.fit_transform(X_categorical[feature].astype(str))
        label_encoders[feature] = le
    
    # 合并所有特征
    X = pd.concat([X_numeric, X_categorical], axis=1)
    
    return X, label_encoders

# 模型训练与可视化
def train_and_visualize():
    # 加载数据
    data = load_data()
    
    # 特征工程
    X, label_encoders = feature_engineering(data)
    y = data['target']
    
    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    # 训练决策树（深度为4）
    dt = DecisionTreeClassifier(max_depth=4, random_state=42)
    dt.fit(X_train, y_train)
    
    # 输出模型准确率
    print(f"模型准确率: {dt.score(X_test, y_test):.4f}")
    
    # 文本形式可视化决策树
    tree_rules = export_text(dt, feature_names=list(X.columns))
    print("\n决策树规则（文本形式）:")
    print(tree_rules)
    
    # 保存文本形式的决策树到文件
    with open('decision_tree_rules.txt', 'w', encoding='utf-8') as f:
        f.write(tree_rules)
    print("\n决策树规则已保存到 decision_tree_rules.txt 文件")
    
    # 图形化可视化决策树
    plt.figure(figsize=(20, 10))
    plot_tree(dt, feature_names=X.columns, class_names=['不会达到100万+', '会达到100万+'], 
              filled=True, rounded=True, fontsize=10)
    plt.title('客户资产达到100万+预测决策树 (深度=4)')
    plt.tight_layout()
    plt.savefig('image_show/decision_tree_visualization_v2.png', dpi=300, bbox_inches='tight')
    plt.close()
    print("\n决策树图形已保存到 decision_tree_visualization_v2.png 文件")
    
    # 显示特征重要性
    feature_importance = pd.DataFrame({
        'feature': X.columns,
        'importance': dt.feature_importances_
    }).sort_values('importance', ascending=False)
    
    print("\n特征重要性排序:")
    print(feature_importance.head(10))
    
    # 保存特征重要性到CSV文件
    feature_importance.to_csv('image_show/decision_tree_feature_importance.csv', index=False)
    print("\n特征重要性已保存到 decision_tree_feature_importance.csv 文件")

if __name__ == '__main__':
    train_and_visualize()